#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
import time
import json
import struct
import shutil
import traceback
from pathlib import Path
from typing import Dict, Any, Optional, Tuple, Iterable
import threading                 # === 新增/保留：线程基元 ===
import queue                     # === 新增/保留：有界队列 ===
import torch
import argparse
# ===== 标准依赖（图像/数值） =====
import cv2
import numpy as np
from PIL import Image
from multiprocessing import Manager
import socket

# ===== 你原有的导入（保持） =====
from gears.netgear_udp import NetGearUDP as NetGear

# ===== 绑定地址与端口 =====
BIND_ADDR = "114.212.86.152"  # 127.0.0.1, 114.212.86.152
PORT = 5558

# ===== 回传给 client 的地址 =====
CLIENT_ADDR = "10.22.36.246"  # "172.27.144.86" "127.0.0.1""172.27.155.100"
CLIENT_PORT = 5559

# # ===== 保存目录（保持你的默认路径；首次运行清空重建） =====
# DATA_DIR = Path("/data/wxk/workspace/mirage/dataset/video000/Viduce/sender_out")
# SAVE_DIR = Path("/data/wxk/workspace/mirage/dataset/video000/Viduce/reciever_out")
# if SAVE_DIR.exists():
#     shutil.rmtree(SAVE_DIR)
# SAVE_DIR.mkdir(parents=True, exist_ok=True)


# ===== NS-3 收包函数（保持）=====
try:
    from ns3.receiver import receive_packet
except Exception as e:
    raise RuntimeError(f"[Init] 无法导入 ns3 函数 receive_packet: {e}")

# ===== ACK/RES 头部格式（保持）=====
_ACK_FMT = "!id"
_ACK_SIZE = struct.calcsize(_ACK_FMT)
RES_MAGIC = b'RSID'
RES_HDR_FMT = '!4sI'  # magic(4s) + frame_id(u32)



# --- 全局对象（带注释） ---
_q_decode = queue.Queue(maxsize=16)  # 有界队列：待解码的 frame_id；背压上游
_q_detect = queue.Queue(maxsize=16)  # 有界队列：待 YOLO 的 (frame_id, png_path, dec_ms)
_stop_event = threading.Event()      # 结束信号



# -------------------------------
# 发送权重文件接口
# -------------------------------
def send_weights(file_path: str, target_ip: str, target_port: int, delay: float = 0.001):
    """
    通过 TCP 将权重文件发送到目标 IP 和端口。
    支持分块发送并在块之间添加延时以避免网络拥塞。
    :param file_path: 要发送的文件路径
    :param target_ip: 目标机器 IP 地址
    :param target_port: 目标端口号
    :param delay: 每发送一个块后暂停的秒数（默认 1 毫秒）
    """
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"文件不存在: {file_path}")

    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        print(f"🔗 正在连接 {target_ip}:{target_port} ...")
        s.connect((target_ip, target_port))

        file_size = os.path.getsize(file_path)
        file_name = os.path.basename(file_path)

        # 发送头部：文件名|大小
        header = f"{file_name}|{file_size}".encode()
        s.sendall(header + b"\n")
        print(f"📤 开始发送文件: {file_name} ({file_size} 字节)")

        # 逐块发送文件内容，带发送间隔
        sent_bytes = 0
        start_time = time.time()

        with open(file_path, "rb") as f:
            while chunk := f.read(8388608):
                s.sendall(chunk)
                sent_bytes += len(chunk)

                # 控制发送间隔，避免网络拥塞
                # if delay > 0:
                #     time.sleep(delay)

        end_time = time.time()
        duration = end_time - start_time
        print(f"✅ 文件发送完成 ({sent_bytes} 字节, 用时 {duration:.2f} 秒)")


def send_weights_udp(file_path: str, target_ip: str, target_port: int, delay: float = 0.001):
    """
    使用 UDP 发送权重文件。
    文件会被分块，每个块附带编号以便接收方重组。
    """
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"文件不存在: {file_path}")

    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    print(f"📡 正在向 {target_ip}:{target_port} 发送文件...")

    file_size = os.path.getsize(file_path)
    file_name = os.path.basename(file_path)

    # 发送文件头：file_name|file_size
    header = f"HEADER|{file_name}|{file_size}".encode()
    s.sendto(header, (target_ip, target_port))
    print(f"📤 文件头已发送: {file_name} ({file_size} 字节)")

    # 每个 UDP 数据包最大安全尺寸（建议 < 60KB）
    CHUNK_SIZE = 60 * 1024  

    sent_bytes = 0
    seq = 0
    start_time = time.time()

    with open(file_path, "rb") as f:
        while chunk := f.read(CHUNK_SIZE):
            packet = f"DATA|{seq}|".encode() + chunk
            s.sendto(packet, (target_ip, target_port))
            sent_bytes += len(chunk)
            seq += 1

            if delay > 0:
                time.sleep(delay)

    # 发送结束标志
    s.sendto(b"END", (target_ip, target_port))

    end_time = time.time()
    s.close()

    duration = end_time - start_time
    print(f"✅ 文件发送完成 ({sent_bytes} 字节, 用时 {duration:.2f} 秒, 共 {seq} 个包)")

# 定义一个锁，防止多个线程同时发送相同文件（可选）
send_lock = threading.Lock()

def async_send_weights():
    """独立线程执行文件发送"""
    # 防止重复并发发送
    if send_lock.locked():
        print("⏳ 上一个文件还在发送中，跳过本次发送。")
        return

    def _send():
        with send_lock:
            try:
                print("📤 [异步线程] 开始发送权重文件...")
                send_weights("/data/shb/test/latest_deltas.pth", CLIENT_ADDR, 5555)  # 每 100 帧发送一次权重文件
                # send_weights_udp("/data/shb/test/latest_deltas.pth", CLIENT_ADDR, 5555)
                print("✅ [异步线程] 文件发送完成。")
            except Exception as e:
                print(f"⚠️ [异步线程] 文件发送出错: {e}")

    # 启动后台线程执行
    threading.Thread(target=_send, daemon=True).start()

def main():
    """
    主流程：
      - 接收 NS-3 送来的 UDP 包
      - receive_packet(data) → “刚刚完整的帧号列表”
      - 对每个新完成的帧 push -> _q_decode（阻塞时形成背压）
      - 两个工作线程常驻：_decoder_worker, _yolo_worker
      - 【新增】在此处一次性加载 YOLOContext，并传入 YOLO 工作线程
    """
    parser = argparse.ArgumentParser()
    parser.add_argument("--gcc", action="store_true",
                        help="enable GCC at receiver side (REMB/arrival-based control)")
    args = parser.parse_args()
    manager = Manager()
    done_frames_shared = manager.list()

    net = NetGear(
        address=BIND_ADDR,
        port=PORT,
        protocol="udp",
        receive_mode=True,
        logging=True,
        mtu=1500,
        recv_buffer_size=32 * 1024 * 1024,
        send_buffer_size=32 * 1024 * 1024,
        queue_maxlen=655360,
        slow_log_dir="./tmp/logs_receiver",          # 这端一般不会写 slow 日志，但可保留
        slow_weights_path="./tmp/slow_module_weights.json",
        rtcp_interval_ms=100,
    )
    net._peer_addr = (CLIENT_ADDR, CLIENT_PORT)
    net.enable_gcc(True)

    # === 在实例创建后，从实例上“绑定” pkt 常量到本地变量 ===
    # 这样后面代码仍可以用 PKT_DATA 等名字，不改流程。
    global PKT_DATA, PKT_SV_DATA, PKT_RES, PKT_TERM, PKT_ACK_FRAME, PKT_ACK_PACKET, PKT_PING, PKT_PONG
    PKT_DATA = net.PKT_DATA
    PKT_SV_DATA = net.PKT_SV_DATA
    PKT_RES = net.PKT_RES
    PKT_TERM = net.PKT_TERM
    PKT_ACK_FRAME = net.PKT_ACK_FRAME
    PKT_ACK_PACKET = net.PKT_ACK_PACKET
    PKT_PING = net.PKT_PING
    PKT_PONG = net.PKT_PONG
    seen = {}  # type: dict[int, float]
    # [新增变量] last_cleanup_ts：上次清理的时间戳，用于控制清理频率（每 10 秒一次）
    last_cleanup_ts = time.time()
    new_frame_count = 0  # 自上次发送以来累计的新帧数量
    SEND_INTERVAL = 1000  # 每收到多少新帧发送一次权重文件  
    try:
        while True:
            pkt = net.recv()
            if pkt is not None:
                ptype = pkt.get("pkt_type")
                data = pkt.get("data", b"")
                if ptype in {PKT_DATA, PKT_SV_DATA, PKT_RES}:
                    frame_ids, packet_id = receive_packet(data)
                    if packet_id:
                        now = time.time()
                        ack_payload = struct.pack(_ACK_FMT, int(packet_id), now)
                        net.send(ack_payload, pkt_type=PKT_ACK_PACKET)

                    # 对“刚刚完整”的帧逐一去重并 ACK（帧级 ACK）
                    if frame_ids:
                        now = time.time()
                        for fid in frame_ids:
                            # 只对未见过的帧做 ACK，并记录时间戳
                            if fid not in seen:
                                # print(f"receive frame {fid}")
                                seen[fid] = now  # 记录首次见到该帧的时间戳
                                new_frame_count += 1  # 记录新增帧
                                # print(f"[Receiver] 收到新帧 {fid}，累计新帧数：{new_frame_count}")  
                                            # 到达阈值就触发发送
                                if new_frame_count >= SEND_INTERVAL:
                                    # print(f"[Receiver] 已累计 {new_frame_count} 新帧，触发发送。")
                                    async_send_weights()
                                    new_frame_count = 0  # 重置计数
                                ack_payload = struct.pack(_ACK_FMT, int(fid), now)
                                net.send(ack_payload, pkt_type=PKT_ACK_FRAME)

                # ===== 新增：每 10 秒清理一次 seen 中的过期项 =====
                # 功能：控制内存占用，把 10 秒前记录的 frame_id 移除（这些帧已完成 ACK、无需再保留）
                now = time.time()
                if now - last_cleanup_ts >= 10.0:
                    expire_before = now - 10.0
                    # 计算需要删除的键列表（避免边遍历边删除）
                    to_delete = [fid for fid, ts in seen.items() if ts < expire_before]
                    if to_delete:
                        for fid in to_delete:
                            seen.pop(fid, None)
                    last_cleanup_ts = now
                # ===== 新增清理结束 =====
                continue

            time.sleep(0.001)

    except KeyboardInterrupt:
        pass
    finally:
        _stop_event.set()
        # 尝试等队列清空（可选）
        try:
            _q_decode.join()
            _q_detect.join()
        except Exception:
            pass
        if hasattr(net, "get_stats"):
            print(f"[Server] Receiver stats: {net.get_stats()}")
        net.close()




if __name__ == "__main__":
    # async_send_weights()
    main()
